import numpy as np
import matplotlib.pyplot as plt


class SimpleLinearRegression2:
    def __init__(self):
        self.a = None
        self.b = None

    def fit(self, xTrain, yTrain):
        assert xTrain.ndim == 1, \
            'Simple Linear Regressor can only solve single feature training data.'

        assert len(xTrain) == len(yTrain), \
            'the size of xTrain must be equal to the size of yTrain'

        xMean = np.mean(xTrain)
        yMean = np.mean(yTrain)

        num = 0.0
        d = 0.0

        num = (xTrain - xMean).dot(yTrain - yMean)
        d = (xTrain - xMean).dot(xTrain - xMean)

        # for x, y in zip(xTrain, yTrain):
        #     num += (x - xMean) * (y - yMean)
        #     d += (x - xMean) ** 2

        self.a = num / d
        self.b = yMean - self.a * xMean

        return self

    def predict(self, xPredict):
        assert xPredict.ndim == 1, \
            'Simple Linear Regressor can only solve single feature training data.'

        assert self.a is not None and self.b is not None, \
            'must fit before predict!'

        return np.array([self._predict(x) for x in xPredict])

    def _predict(self, xSingle):
        return self.a * xSingle + self.b

    def __repr__(self):
        return "SimpleLinearRegression2()"


if __name__ == '__main__':
    x = np.array([1., 2., 3., 4., 5.])
    y = np.array([1., 3., 2., 3., 5.])

    simpleLinearRegression = SimpleLinearRegression2()
    simpleLinearRegression.fit(x, y)
    xPredict = 6
    predict = simpleLinearRegression.predict(np.array([xPredict]))
    print(predict, simpleLinearRegression.a, simpleLinearRegression.b)
    yHat = simpleLinearRegression.predict(x)

    plt.scatter(x, y)
    plt.plot(x, yHat, color='r')
    plt.axis([0, 6, 0, 6])
    plt.show()

